# 改进版首板策略
# 首板策略: 判断尾盘涨幅超过5%
# 如果过去1天涨幅超过5% 3天涨幅超过8% 10天超15% 30天超20%取消买入，静默15天
# 根据其他参数（缩量上涨还要涨、双均线、利好利空、北向资金）决定买入仓位，尾盘买入
# 买入后，第二天5%止损，10%止盈，最大持有时间5天
# 买入后，当日最大跌幅超3%，第二天首盘卖出
# 买入后，当日上涨，次日最高价浮动3%止盈
from datacache import get_stock_data
import matplotlib.pyplot as plt
import pandas as pd

def main():
    stock_code = 'sh600185'  # 平安银行
    start_date = '20200101'
    end_date = '20250130'
    buy_raise_limit = 0.08 # 买入涨幅限制

    wallet_funds = 10000 # 初始资金10000
    holding = 0 # 持有股票数量
    sell_warning = False # 卖出警告 下一个交易日开盘卖出
    holding_price = 0 # 持有股票价格
    sum_profit = [] # 收益曲线

    df = get_stock_data(stock_code, start_date, end_date)
    for index, row in df.iterrows():
        row_date = row['date']
        close_price = row['close']
        open_price = row['open']
        high_price = row['high'] # 最高价
        low_price = row['low'] # 最低价
        volume = row['volume']
        turnover = row['turnover'] # 换手率
        volume_rate = 0 # 5日量比
        if index > 5:
            value1 = df['volume'].values[index - 5]
            value2 = df['volume'].values[index - 4]
            value3 = df['volume'].values[index - 3]
            value4 = df['volume'].values[index - 2]
            value5 = df['volume'].values[index - 1]
            average_volumn = (value1 + value2 + value3 + value4 + value5) / 6
            volume_rate = volume / average_volumn - 1
        amplitude_0 = (close_price - open_price) / open_price # 今日涨跌幅度
        if amplitude_0 > buy_raise_limit and holding == 0:
            # 触发买入涨幅条件，且未持有股票准备买入
            amplitude_1 = 0 # 1日涨跌
            if index > 0:
                amplitude_1 = (close_price - df.loc[index-1, 'close']) / df.loc[index-1, 'close']
            amplitude_3 = 0 # 3日涨跌
            if index > 2:
                amplitude_3 = (close_price - df.loc[index-3, 'close']) / df.loc[index-3, 'close']
            amplitude_10 = 0 # 10日涨跌
            if index > 9:
                amplitude_10 = (close_price - df.loc[index-10, 'close']) / df.loc[index-10, 'close']
            amplitude_30 = 0 # 30日涨跌
            if index > 29:
                amplitude_30 = (close_price - df.loc[index-30, 'close']) / df.loc[index-30, 'close']
            if amplitude_1 > 0.05 and amplitude_3 > 0.08 and amplitude_10 > 0.15 and amplitude_30 > 0.20:
                print("取消买入了")
                pass # 取消买入
            position = 1 # 买入仓位
            prepared_count = wallet_funds * position / close_price # 准备买入数量
            wallet_funds = wallet_funds - prepared_count * close_price # 减去买入金额
            holding = prepared_count # 更新持有数量
            holding_price = close_price # 更新持有价格
        else:
            # 持有股票，准备卖出
            if sell_warning: # 卖出预警，开盘直接卖
                prepared_amount = open_price * holding
                wallet_funds = wallet_funds + prepared_amount # 加上卖出金额
                holding = 0 # 更新持有数量
                holding_price = close_price  # 更新持有价格
                pass # 卖出
            amplitude_min = (open_price - low_price) / open_price # 当日最大跌幅
            amplitude_max = (high_price - open_price) / open_price  # 当日最大涨幅
            amplitude_base_min = (low_price - holding_price) / open_price # 以成本价格为基准的最大涨幅
            print(amplitude_base_min)
            if amplitude_min > 0.05 and volume_rate < -1:
                if amplitude_base_min < 0.085:
                    # 跌幅超过 5%，5%止损 并且缩量 并且基于成本价的收益小于了8.5%
                    prepared_amount = open_price * holding * 0.95
                    wallet_funds = wallet_funds + prepared_amount # 加上卖出金额
                    holding = 0  # 更新持有数量
                    holding_price = close_price  # 更新持有价格
                    pass
            if amplitude_0 < 0:
                sell_warning = True # 当日收益为负，次日止损
            pass
        sum_profit.append(wallet_funds + holding * holding_price)
    print("总资产" + str(wallet_funds + holding * holding_price))
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
    plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(14, 7), sharex=True)
    df['date'] = pd.to_datetime(df['date'])
    # 获取日期部分
    df['date'] = df['date'].dt.date
    # 绘制价格和移动平均线
    ax1.plot(df['date'], df['close'], label='收盘价', color='#0d0dff80')
    ax1.plot(df['date'], df['ma5'], label='5日均线', color='red')
    ax1.plot(df['date'], df['ma20'], label='20日均线', color='green')

    ax1.set_title('股票代码:' + stock_code)
    ax1.set_ylabel('价格')
    ax1.legend()

    ax2.set_title('总资产:' + stock_code)
    ax2.plot(df['date'], sum_profit, label='日期', color='red')
    ax2.set_ylabel('总量')
    ax2.legend()
    plt.xlabel('日期')
    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    main()